#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
JWT认证中间件
"""

import re
from fastapi import Request, HTTPException, status
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from typing import Optional, Dict, Any

from src.utils.token_generator import token_generator
from src.utils.response import ApiResponse
from src.constants.resultcode import ResultCode
from src.utils.logger import get_logger
import config

logger = get_logger("capcut")


class JWTAuthMiddleware(BaseHTTPMiddleware):
    """JWT认证中间件"""
    
    def __init__(self, app):
        super().__init__(app)
        self.white_list = config.JWT_WHITE_LIST
        self.header_name = config.JWT_HEADER_NAME
        
        # 编译白名单正则表达式（支持通配符）
        self.white_list_patterns = []
        for path in self.white_list:
            # 转换通配符为正则表达式
            pattern = path.replace('*', '.*').replace('?', '.')
            self.white_list_patterns.append(re.compile(f"^{pattern}$"))
    
    def is_path_in_whitelist(self, path: str) -> bool:
        """检查路径是否在白名单中"""
        # 精确匹配
        if path in self.white_list:
            return True
        
        # 正则匹配（支持通配符）
        for pattern in self.white_list_patterns:
            if pattern.match(path):
                return True
        
        return False
    
    def extract_token(self, request: Request) -> Optional[str]:
        """从请求中提取Token"""
        # 从自定义头部提取
        print(request.headers)
        token = request.headers.get(self.header_name)
        if token:
            logger.info(token)
            return token
        
        logger.info(f"未找到Token，检查的头部: {self.header_name}")
        return None
    
    def create_bearer_token(self, token: str) -> str:
        """为Token添加Bearer前缀"""
        return f"Bearer {token}"
    
    async def dispatch(self, request: Request, call_next):
        """中间件主要逻辑"""
        path = request.url.path
        method = request.method
        
        # 记录请求
        logger.info(f"{method} {path} - 来自 {request.client.host}")
        
        # 检查是否在白名单中
        if self.is_path_in_whitelist(path):
            logger.info(f"路径在白名单中: {path}")
            response = await call_next(request)
            return response
        
        # 提取Token
        token = self.extract_token(request)
        
        if not token:
            logger.warning(f"缺少认证Token: {method} {path}")
            logger.debug(f"请求头部: {dict(request.headers)}")
            return ApiResponse.error(
                result_code=ResultCode.AUTHENTICATION_FAILED,
                message=f"缺少认证Token，请在请求头中包含 {self.header_name} 字段",
                status_code=status.HTTP_401_UNAUTHORIZED
            )
        
        # 验证Token
        try:
            payload = token_generator.verify_token(token)
            if not payload:
                logger.warning(f"无效的Token: {token[:20]}...")
                return ApiResponse.error(
                    result_code=ResultCode.AUTHENTICATION_FAILED,
                    message="无效的Token",
                    status_code=status.HTTP_401_UNAUTHORIZED
                )
            
            # 检查用户状态（可选，需要查询数据库）
            user_id = payload.get("user_id")
            if user_id:
                # 这里可以添加用户状态检查逻辑
                # user_status = await check_user_status(user_id)
                # if user_status != 1:
                #     return unauthorized_response("用户已被禁用")
                pass
            
            # 将用户信息添加到请求状态中
            request.state.user_id = payload.get("user_id")
            request.state.username = payload.get("username")
            request.state.token_id = payload.get("jti")
            request.state.token_payload = payload
            
            # 为后续处理添加Bearer格式的Token
            bearer_token = self.create_bearer_token(token)
            request.state.bearer_token = bearer_token
            
            logger.debug(f"Token验证成功: 用户 {payload.get('username')} (ID: {user_id})")
            
        except Exception as e:
            logger.error(f"Token验证异常: {str(e)}", exc_info=True)
            return ApiResponse.error(
                result_code=ResultCode.AUTHENTICATION_FAILED,
                message="Token验证失败",
                status_code=status.HTTP_401_UNAUTHORIZED
            )
        
        # 继续处理请求
        response = await call_next(request)
        
        # 添加响应头（可选）
        response.headers["X-Token-Valid"] = "true"
        
        return response


# JWT异常类
class JWTException(HTTPException):
    """JWT相关异常"""
    
    def __init__(self, detail: str = "JWT认证失败", status_code: int = status.HTTP_401_UNAUTHORIZED):
        super().__init__(status_code=status_code, detail=detail)


# 工具函数
def get_current_user_from_request(request: Request) -> Optional[Dict[str, Any]]:
    """从请求中获取当前用户信息"""
    if hasattr(request.state, 'token_payload'):
        return {
            'user_id': request.state.user_id,
            'username': request.state.username,
            'token_id': request.state.token_id,
            'payload': request.state.token_payload
        }
    return None


 